#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
from pathlib import Path
from PIL import Image, ImageOps
import argparse
from typing import List, Tuple, Optional
import shutil

class ImageCompressor:
    def __init__(self):
        self.supported_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
        self.max_images = 100
        self.processed_count = 0
    
    def is_image(self, file_path: Path) -> bool:
        return file_path.suffix.lower() in self.supported_formats
    
    def get_unique_name(self, base_path: Path, suffix: str = "_cp") -> Path:
        stem = base_path.stem
        parent = base_path.parent
        extension = base_path.suffix
        
        new_name = f"{stem}{suffix}{extension}"
        new_path = parent / new_name
        
        counter = 1
        while new_path.exists():
            new_name = f"{stem}{suffix}_{counter}{extension}"
            new_path = parent / new_name
            counter += 1
        
        return new_path
    
    def get_smart_quality(self, file_path: Path, base_quality: int) -> int:
        # 尊重用户设定的质量，只对小文件做适度提升优化
        file_size = file_path.stat().st_size / 1024  # KB
        
        # 对于小文件，可以适度提升质量以获得更好效果
        if file_size < 100 and base_quality < 80:  # 小于100KB且用户设置质量低
            return min(85, base_quality + 15)  # 适度提升
        elif file_size < 500 and base_quality < 70:  # 100KB-500KB且质量低  
            return min(80, base_quality + 10)   # 适度提升
        else:
            # 其他情况一律使用用户设定值
            return base_quality

    def compress_to_target_size(self, img: Image.Image, output_path: Path, 
                               target_size_kb: int, format_change: Optional[str], 
                               initial_quality: int = 85, original_size_kb: float = 0) -> Tuple[bool, int, float]:
        """迭代压缩直到满足目标大小要求"""
        import tempfile
        
        # 严格按照用户设定的质量进行压缩
        current_quality = initial_quality
        min_quality = max(5, initial_quality - 10)  # 最低质量限制
        step_size = 5  # 精细控制步长
        
        # 如果没有指定格式且目标很小，自动转换为JPEG获得更好压缩
        auto_format = False
        if not format_change and target_size_kb <= 500:
            format_change = 'jpeg'
            auto_format = True
            if img.mode == 'RGBA':
                background = Image.new('RGB', img.size, (255, 255, 255))
                background.paste(img, mask=img.split()[-1])
                img = background
        
        # 先试试初始质量能否满足要求
        with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{format_change or "jpg"}') as temp_file:
            temp_path = Path(temp_file.name)
        
        try:
            save_kwargs = {'optimize': True}
            if format_change and format_change.lower() in ['jpeg', 'jpg']:
                save_kwargs['quality'] = current_quality
            elif format_change and format_change.lower() == 'webp':
                save_kwargs['quality'] = current_quality
            elif format_change and format_change.lower() == 'png':
                # PNG 使用 compress_level 而不是 quality
                compress_level = min(9, max(0, int((100 - current_quality) / 10)))
                save_kwargs['compress_level'] = compress_level
            
            img.save(temp_path, format=format_change.upper() if format_change else 'JPEG', **save_kwargs)
            temp_size_kb = temp_path.stat().st_size / 1024
            
            # 如果初始质量就满足要求且不比原文件大，直接返回
            if temp_size_kb <= target_size_kb and (original_size_kb == 0 or temp_size_kb <= original_size_kb):
                shutil.move(str(temp_path), str(output_path))
                return True, current_quality, temp_size_kb
            else:
                temp_path.unlink()
        except Exception as e:
            if temp_path.exists():
                temp_path.unlink()
            return False, current_quality, 0
        
        # 需要进一步压缩，逐步降低质量
        while current_quality >= min_quality:
            # 先降低质量
            current_quality -= step_size
            
            # 确保不低于最低限制
            if current_quality < min_quality:
                current_quality = min_quality
            
            with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{format_change or "jpg"}') as temp_file:
                temp_path = Path(temp_file.name)
            
            try:
                save_kwargs = {'optimize': True}
                if format_change and format_change.lower() in ['jpeg', 'jpg']:
                    save_kwargs['quality'] = current_quality
                elif format_change and format_change.lower() == 'webp':
                    save_kwargs['quality'] = current_quality
                elif format_change and format_change.lower() == 'png':
                    # PNG 使用 compress_level 而不是 quality
                    compress_level = min(9, max(0, int((100 - current_quality) / 10)))
                    save_kwargs['compress_level'] = compress_level
                
                img.save(temp_path, format=format_change.upper() if format_change else 'JPEG', **save_kwargs)
                temp_size_kb = temp_path.stat().st_size / 1024
                
                if temp_size_kb <= target_size_kb and (original_size_kb == 0 or temp_size_kb <= original_size_kb):
                    shutil.move(str(temp_path), str(output_path))
                    return True, current_quality, temp_size_kb
                else:
                    temp_path.unlink()
                    
                # 如果已经是最低质量还不满足，跳出循环
                if current_quality <= min_quality:
                    break
                    
            except Exception as e:
                if temp_path.exists():
                    temp_path.unlink()
                return False, current_quality, 0
        
        # 如果最低质量仍不满足，使用最低质量保存
        try:
            save_kwargs = {'optimize': True}
            if format_change and format_change.lower() in ['jpeg', 'jpg']:
                save_kwargs['quality'] = min_quality
            elif format_change and format_change.lower() == 'webp':
                save_kwargs['quality'] = min_quality
            elif format_change and format_change.lower() == 'png':
                # PNG 使用 compress_level 而不是 quality
                compress_level = min(9, max(0, int((100 - min_quality) / 10)))
                save_kwargs['compress_level'] = compress_level
            
            img.save(output_path, format=format_change.upper() if format_change else 'JPEG', **save_kwargs)
            final_size_kb = output_path.stat().st_size / 1024
            
            # 如果压缩后文件比原文件还大，且没有真正的格式转换，返回失败让上层处理
            if original_size_kb > 0 and final_size_kb > original_size_kb:
                # 如果没有真正的格式转换，则返回失败让上层处理
                if not format_change:
                    if output_path.exists():
                        output_path.unlink()
                    return False, min_quality, final_size_kb
                
            return True, min_quality, final_size_kb
        except Exception:
            return False, min_quality, 0

    def compress_image(self, input_path: Path, output_path: Path, quality: int = 85, 
                      format_change: Optional[str] = None, target_size_kb: int = 500) -> bool:
        try:
            file_size_kb = input_path.stat().st_size / 1024
            
            # 小文件直接复制（可配置阈值，默认100KB）
            if file_size_kb < 100:
                shutil.copy2(input_path, output_path)
                print(f"文件过小({file_size_kb:.1f}KB)，直接复制: {output_path}")
                return True
            
            # 如果原文件已经满足目标大小且质量较好，直接复制
            if file_size_kb <= target_size_kb and not format_change:
                shutil.copy2(input_path, output_path)
                print(f"文件已符合要求({file_size_kb:.1f}KB≤{target_size_kb}KB)，直接复制: {output_path}")
                return True
            
            with Image.open(input_path) as img:
                img = ImageOps.exif_transpose(img)
                
                if format_change:
                    output_path = output_path.with_suffix(f'.{format_change.lower()}')
                
                # 尝试迭代压缩到目标大小
                success, final_quality, output_size_kb = self.compress_to_target_size(
                    img, output_path, target_size_kb, format_change, quality, file_size_kb
                )
                
                if success:
                    # 检查是否真的有格式转换
                    original_ext = input_path.suffix.lower().lstrip('.')
                    # 将jpeg和jpg视为相同格式
                    if original_ext == 'jpeg':
                        original_ext = 'jpg'
                    target_ext = format_change.lower() if format_change else original_ext
                    if target_ext == 'jpeg':
                        target_ext = 'jpg'
                    is_real_format_change = format_change and original_ext != target_ext
                    
                    # 如果压缩后文件比原文件还大，且没有真正的格式转换，直接复制原文件
                    if output_size_kb > file_size_kb and not is_real_format_change:
                        # 删除已生成的压缩文件
                        if output_path.exists():
                            output_path.unlink()
                        shutil.copy2(input_path, output_path)
                        print(f"压缩后文件更大，直接复制原文件: {output_path} ({file_size_kb:.1f}KB)")
                        return True
                    elif output_size_kb > file_size_kb and is_real_format_change:
                        # 真正的格式转换导致文件变大时，给出警告但仍保存
                        print(f"格式转换完成: {output_path} (⚠ {file_size_kb:.1f}KB→{output_size_kb:.1f}KB, 质量:{final_quality}% 格式转换导致文件变大)")
                        return True
                    
                    compression_ratio = (1 - output_size_kb / file_size_kb) * 100
                    if output_size_kb <= target_size_kb:
                        status = "✓"
                        size_info = ""
                    else:
                        status = "⚠"
                        size_info = f" 超出目标{target_size_kb}KB"
                    print(f"压缩完成: {output_path} ({status} {file_size_kb:.1f}KB→{output_size_kb:.1f}KB, 质量:{final_quality}%{size_info})")
                    return True
                else:
                    print(f"压缩失败 {input_path}")
                    return False
                
        except Exception as e:
            print(f"压缩失败 {input_path}: {e}")
            return False
    
    def compress_single_image(self, image_path: str, quality: int = 85, 
                            format_change: Optional[str] = None, target_size_kb: int = 500) -> bool:
        input_path = Path(image_path)
        if not input_path.exists() or not self.is_image(input_path):
            print(f"文件不存在或不是支持的图片格式: {image_path}")
            return False
        
        output_path = self.get_unique_name(input_path)
        if format_change:
            output_path = output_path.with_suffix(f'.{format_change.lower()}')
        
        success = self.compress_image(input_path, output_path, quality, format_change, target_size_kb)
        if success:
            print(f"压缩完成: {output_path}")
        return success
    
    def get_unique_folder_name(self, base_path: Path, folder_name: str = "compress") -> Path:
        folder_path = base_path / folder_name
        counter = 1
        
        while folder_path.exists():
            folder_path = base_path / f"{folder_name}_{counter}"
            counter += 1
        
        return folder_path
    
    def compress_batch_images(self, image_paths: List[str], quality: int = 85, 
                            format_change: Optional[str] = None, target_size_kb: int = 500) -> bool:
        if len(image_paths) > self.max_images:
            print(f"图片数量超过限制 ({self.max_images}张)")
            return False
        
        if not image_paths:
            print("没有找到图片文件")
            return False
        
        base_dir = Path(image_paths[0]).parent
        compress_folder = self.get_unique_folder_name(base_dir)
        compress_folder.mkdir(exist_ok=True)
        
        success_count = 0
        for image_path in image_paths:
            input_path = Path(image_path)
            if not self.is_image(input_path):
                continue
            
            filename = input_path.name
            if format_change:
                filename = input_path.stem + f".{format_change.lower()}"
            
            output_path = compress_folder / filename
            
            if self.compress_image(input_path, output_path, quality, format_change, target_size_kb):
                success_count += 1
                print(f"压缩完成: {output_path}")
        
        print(f"批量压缩完成，成功压缩 {success_count} 张图片到 {compress_folder}")
        return success_count > 0
    
    def compress_folder_recursive(self, folder_path: str, quality: int = 85, 
                                format_change: Optional[str] = None, target_size_kb: int = 500) -> bool:
        folder = Path(folder_path)
        if not folder.exists() or not folder.is_dir():
            print(f"文件夹不存在: {folder_path}")
            return False
        
        self.processed_count = 0
        success_count = self._compress_folder_recursive_helper(folder, folder, quality, format_change, target_size_kb)
        
        if success_count > 0:
            print(f"递归压缩完成，总共压缩 {success_count} 张图片")
        else:
            print("没有找到可压缩的图片")
        
        return success_count > 0
    
    def _compress_folder_recursive_helper(self, current_folder: Path, root_folder: Path, 
                                        quality: int, format_change: Optional[str], target_size_kb: int) -> int:
        success_count = 0
        
        for item in current_folder.iterdir():
            if self.processed_count >= self.max_images:
                print(f"已达到最大处理数量限制 ({self.max_images}张)")
                break
            
            if item.is_file() and self.is_image(item):
                relative_path = item.relative_to(root_folder)
                
                compressed_name = item.name
                if format_change:
                    compressed_name = item.stem + f".{format_change.lower()}"
                
                output_path = item.parent / compressed_name
                
                if self.compress_image(item, output_path, quality, format_change, target_size_kb):
                    success_count += 1
                    self.processed_count += 1
                    print(f"压缩完成: {output_path}")
            
            elif item.is_dir():
                success_count += self._compress_folder_recursive_helper(item, root_folder, quality, format_change, target_size_kb)
        
        return success_count

def main():
    parser = argparse.ArgumentParser(description='图片压缩工具')
    parser.add_argument('path', help='图片文件路径或文件夹路径')
    parser.add_argument('--mode', choices=['single', 'batch', 'folder'], default='single',
                       help='压缩模式: single(单张), batch(批量), folder(文件夹递归)')
    parser.add_argument('--quality', type=int, default=85, choices=range(1, 101),
                       help='压缩质量 (1-100, 默认85)')
    parser.add_argument('--format', choices=['jpeg', 'png', 'webp'], 
                       help='转换格式 (可选)')
    parser.add_argument('--target-size', type=int, default=500,
                       help='目标文件大小限制(KB, 默认500)')
    
    args = parser.parse_args()
    
    compressor = ImageCompressor()
    
    if args.mode == 'single':
        compressor.compress_single_image(args.path, args.quality, args.format, args.target_size)
    
    elif args.mode == 'batch':
        if os.path.isdir(args.path):
            image_files = []
            for file in Path(args.path).iterdir():
                if file.is_file() and compressor.is_image(file):
                    image_files.append(str(file))
            compressor.compress_batch_images(image_files, args.quality, args.format, args.target_size)
        else:
            print("批量模式需要提供文件夹路径")
    
    elif args.mode == 'folder':
        compressor.compress_folder_recursive(args.path, args.quality, args.format, args.target_size)

def interactive_mode():
    compressor = ImageCompressor()
    
    print("=== 图片压缩工具 ===")
    print("1. 单张图片压缩")
    print("2. 批量图片压缩")
    print("3. 文件夹递归压缩")
    print("4. 退出")
    
    while True:
        choice = input("\n请选择功能 (1-4): ").strip()
        
        if choice == '4':
            print("退出程序")
            break
        
        if choice not in ['1', '2', '3']:
            print("无效选择，请重新输入")
            continue
        
        path = input("请输入文件/文件夹路径: ").strip()
        
        print("\n压缩质量选择:")
        print("1. 高质量低压缩 (质量: 95)")
        print("2. 中等质量中压缩 (质量: 85)")
        print("3. 低质量高压缩 (质量: 60)")
        
        quality_choice = input("请选择压缩质量 (1-3, 默认2): ").strip()
        quality_map = {'1': 95, '2': 85, '3': 60}
        quality = quality_map.get(quality_choice, 85)
        
        target_size_input = input("请输入目标文件大小限制(KB, 默认500): ").strip()
        target_size = int(target_size_input) if target_size_input.isdigit() else 500
        
        format_change = None
        if quality <= 60:
            format_input = input("是否转换为JPEG格式以获得更好的压缩率? (y/n, 默认n): ").strip().lower()
            if format_input == 'y':
                format_change = 'jpeg'
        
        if choice == '1':
            compressor.compress_single_image(path, quality, format_change, target_size)
        elif choice == '2':
            if os.path.isdir(path):
                image_files = []
                for file in Path(path).iterdir():
                    if file.is_file() and compressor.is_image(file):
                        image_files.append(str(file))
                compressor.compress_batch_images(image_files, quality, format_change, target_size)
            else:
                print("批量模式需要提供文件夹路径")
        elif choice == '3':
            compressor.compress_folder_recursive(path, quality, format_change, target_size)

if __name__ == "__main__":
    if len(sys.argv) > 1:
        main()
    else:
        interactive_mode()